import asyncio
import websockets
import random
import json
from datetime import datetime

class MockWebSocketClient:
    def __init__(self, url):
        self.url = url

    async def send_websocket_message(self, user_input, message_id):
        async with websockets.connect(self.url) as websocket:
            print(f"WebSocket connected by {user_input}")
            initial_data = {
                "code": "invite",
                "messageId": message_id,
                "msg": f"{user_input}_0",
                "timestamp": str(datetime.now())
            }
            await websocket.send(json.dumps(initial_data))
            print(f"Initial message  {initial_data}")

            for i in range(1, 20):
                random_role = random.randint(0, 1)
                ack_data = {
                    "code": "ack",
                    "messageId": message_id,
                    "msg": f"{user_input}_{i}",
                    "role": random_role,
                    "timestamp": str(datetime.now())
                }
                await websocket.send(json.dumps(ack_data))

            bye_data = {
                "code": "bye",
                "messageId": message_id,
                "msg": f"{user_input}_0",
                "timestamp": str(datetime.now())
            }
            await websocket.send(json.dumps(bye_data))
            print(f"WebSocket messages sent for {user_input}")

    async def send_messages(self, user_inputs):
        tasks = [self.send_websocket_message(user_input, f"1000{i}") for i, user_input in enumerate(user_inputs)]
        await asyncio.gather(*tasks)

if __name__ == "__main__":
    url = "ws://127.0.0.1:9998/websocket"
    user_inputs = [index for index in range(150)]  # 生成 0 到 39 的 40 个输入
    client = MockWebSocketClient(url)
    asyncio.run(client.send_messages(user_inputs))
